package cn.ycc1.mylangchain4j.controller;

import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

/**
 * @author ycc
 * @date 2025/5/6
 */
@RestController
@RequestMapping("/langchain4j")
public class StreamChatController {
    private final StreamingChatModel streamingChatModel;

    public StreamChatController(StreamingChatModel streamingChatModel) {
        this.streamingChatModel = streamingChatModel;
    }

    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamResponse(@RequestParam String message) {
        return Flux.create(emitter -> {
            streamingChatModel.chat(message, new StreamingChatResponseHandler() {
                @Override
                public void onPartialResponse(String partialResponse) {
                    System.out.println("onPartialResponse: " + partialResponse);
                    emitter.next(partialResponse);
                }

                @Override
                public void onCompleteResponse(ChatResponse completeResponse) {
                    System.out.println("onCompleteResponse: " + completeResponse);
                    // emitter.next("[DONE]"); // 明确发送结束信号（可选）
                    emitter.complete();
                }

                @Override
                public void onError(Throwable error) {
                    error.printStackTrace();
                    emitter.error(error);
                }
            });
        });
    }
}
